"""
HDF5文件处理模块
负责HDF5文件的读取、解析和数据提取
"""

import os
import h5py
import numpy as np
import config


class HDF5FileHandler:
    """HDF5文件处理类，负责所有HDF5文件相关操作"""
    
    def __init__(self):
        self.current_file = None
        self.file_structure = None
        
    def open_file(self, file_path):
        """
        打开HDF5文件
        
        Args:
            file_path: HDF5文件路径
            
        Returns:
            bool: 是否成功打开文件
        """
        try:
            if self.current_file:
                self.current_file.close()
            
            self.current_file = h5py.File(file_path, 'r')
            self.file_structure = self._scan_file_structure()
            return True
        except Exception as e:
            print(f"Error opening HDF5 file: {e}")
            return False
    
    def close_file(self):
        """关闭当前打开的HDF5文件"""
        if self.current_file:
            self.current_file.close()
            self.current_file = None
            self.file_structure = None
    
    def _scan_file_structure(self):
        """
        扫描HDF5文件结构
        
        Returns:
            dict: 文件结构信息
        """
        if not self.current_file:
            return None
        
        structure = {
            'name': '/',
            'type': 'root',
            'children': []
        }
        
        def visit_item(name, obj):
            """递归访问HDF5对象"""
            path_parts = name.split('/')
            current_dict = structure
            
            # 导航到正确的位置
            for part in path_parts[:-1]:
                if part:  # 跳过空字符串
                    # 查找或创建中间组
                    found = False
                    for child in current_dict['children']:
                        if child['name'] == part:
                            current_dict = child
                            found = True
                            break
                    if not found:
                        new_group = {
                            'name': part,
                            'type': 'group',
                            'path': '/'.join(path_parts[:-1]),
                            'children': []
                        }
                        current_dict['children'].append(new_group)
                        current_dict = new_group
            
            # 添加当前对象
            if isinstance(obj, h5py.Group):
                item_info = {
                    'name': path_parts[-1],
                    'type': 'group',
                    'path': name,
                    'children': []
                }
            elif isinstance(obj, h5py.Dataset):
                item_info = {
                    'name': path_parts[-1],
                    'type': 'dataset',
                    'path': name,
                    'shape': obj.shape,
                    'dtype': str(obj.dtype),
                    'size': obj.size,
                    'nbytes': obj.nbytes
                }
            else:
                return  # 跳过其他类型
            
            current_dict['children'].append(item_info)
        
        self.current_file.visititems(visit_item)
        return structure
    
    def get_file_structure(self):
        """获取文件结构"""
        return self.file_structure
    
    def get_dataset_data(self, dataset_path, waveform_index=None):
        """
        获取数据集数据
        
        Args:
            dataset_path: 数据集路径
            waveform_index: 波形索引（对于多维数据）
            
        Returns:
            numpy.ndarray: 数据数组，如果失败返回None
        """
        try:
            if not self.current_file:
                return None
            
            dataset = self.current_file[dataset_path]
            
            if waveform_index is not None and len(dataset.shape) > 1:
                # 对于多维数据，选择特定波形
                return dataset[waveform_index, :]
            else:
                # 对于一维数据或获取全部数据
                return dataset[:]
                
        except Exception as e:
            print(f"Error reading dataset {dataset_path}: {e}")
            return None
    
    def get_dataset_info(self, dataset_path):
        """
        获取数据集信息
        
        Args:
            dataset_path: 数据集路径
            
        Returns:
            dict: 数据集信息
        """
        try:
            if not self.current_file:
                return None
            
            dataset = self.current_file[dataset_path]
            return {
                'shape': dataset.shape,
                'dtype': str(dataset.dtype),
                'size': dataset.size,
                'nbytes': dataset.nbytes,
                'max_waveform_index': dataset.shape[0] - 1 if len(dataset.shape) > 1 else 0
            }
        except Exception as e:
            print(f"Error getting dataset info: {e}")
            return None
    
    def find_amp_phase_pairs(self):
        """
        查找幅度和相位数据对
        
        Returns:
            list: 包含配对信息的列表
        """
        if not self.file_structure:
            return []
        
        pairs = []
        all_datasets = self._get_all_datasets(self.file_structure)
        
        # 按组分类数据集
        groups = {}
        for dataset in all_datasets:
            group_path = '/'.join(dataset['path'].split('/')[:-1])
            if group_path not in groups:
                groups[group_path] = []
            groups[group_path].append(dataset)
        
        # 在每个组中查找配对
        for group_path, datasets in groups.items():
            amp_datasets = []
            phase_datasets = []
            
            for dataset in datasets:
                name = dataset['name'].lower()
                if any(pattern in name for pattern in config.AMP_PHASE_PATTERNS['amp_patterns']):
                    amp_datasets.append(dataset)
                elif any(pattern in name for pattern in config.AMP_PHASE_PATTERNS['phase_patterns']):
                    phase_datasets.append(dataset)
            
            # 尝试配对
            for amp_ds in amp_datasets:
                amp_name = amp_ds['name']
                # 移除amp关键字，寻找对应的phase数据集
                base_name = amp_name
                for pattern in config.AMP_PHASE_PATTERNS['amp_patterns']:
                    if pattern in amp_name:
                        base_name = amp_name.replace(pattern, '')
                        break
                
                # 查找对应的phase数据集
                for phase_ds in phase_datasets:
                    if base_name in phase_ds['name']:
                        pairs.append({
                            'base_name': base_name,
                            'amp_dataset': amp_ds,
                            'phase_dataset': phase_ds,
                            'group': group_path
                        })
                        break
        
        return pairs
    
    def _get_all_datasets(self, structure):
        """
        递归获取所有数据集
        
        Args:
            structure: 文件结构字典
            
        Returns:
            list: 所有数据集的列表
        """
        datasets = []
        
        def collect_datasets(node):
            if node['type'] == 'dataset':
                datasets.append(node)
            elif 'children' in node:
                for child in node['children']:
                    collect_datasets(child)
        
        collect_datasets(structure)
        return datasets
    
    def get_sampling_rate(self):
        """
        从HDF5文件中提取采样率信息
        
        Returns:
            float: 采样率（MHz），如果未找到返回默认值
        """
        try:
            # 尝试从不同可能的位置读取采样率
            sampling_rate_paths = [
                'Group_1/iLinac-CM01:LLRF-01:TriggerWaveformSamplingRate',
                'TriggerWaveformSamplingRate',
                'SamplingRate',
                'sampling_rate'
            ]
            
            for path in sampling_rate_paths:
                try:
                    if path in self.current_file:
                        data = self.current_file[path][:]
                        if len(data) > 0:
                            return float(data[0])
                except:
                    continue
            
            # 如果没有找到，返回默认值
            return config.DEFAULT_SAMPLING_FREQ
            
        except Exception as e:
            print(f"Error getting sampling rate: {e}")
            return config.DEFAULT_SAMPLING_FREQ
    
    def create_preview_text(self, file_path):
        """
        创建文件预览文本
        
        Args:
            file_path: 文件路径
            
        Returns:
            str: 预览文本
        """
        if not self.current_file or not self.file_structure:
            return "无法读取文件信息"
        
        file_size = os.path.getsize(file_path) / (1024 * 1024)  # MB
        
        preview_text = f"文件: {os.path.basename(file_path)}\n"
        preview_text += f"大小: {file_size:.2f} MB\n"
        preview_text += f"采样率: {self.get_sampling_rate()} MHz\n\n"
        
        # 统计组和数据集数量
        groups_count = 0
        datasets_count = 0
        
        def count_items(node):
            nonlocal groups_count, datasets_count
            if node['type'] == 'group':
                groups_count += 1
            elif node['type'] == 'dataset':
                datasets_count += 1
            
            if 'children' in node:
                for child in node['children']:
                    count_items(child)
        
        count_items(self.file_structure)
        
        preview_text += f"包含 {groups_count} 个组，{datasets_count} 个数据集\n"
        
        # 显示配对信息
        pairs = self.find_amp_phase_pairs()
        if pairs:
            preview_text += f"\n找到 {len(pairs)} 个幅度-相位数据对:\n"
            for pair in pairs[:5]:  # 只显示前5个
                preview_text += f"  - {pair['base_name']}\n"
            if len(pairs) > 5:
                preview_text += f"  ... 还有 {len(pairs) - 5} 个\n"
        
        return preview_text 